{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "5BiR_DT_CJqI"
      },
      "source": [
        "##### Copyright 2019 The AdaNet Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "o5UC44mKCO0L"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "DTvBhHYwbRQl"
      },
      "source": [
        "# Customizing AdaNet With TensorFlow Hub Modules"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "tB0jkokjbT0L"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/adanet/blob/master/adanet/examples/tutorials/customizing_adanet_with_tfhub.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/adanet/blob/master/adanet/examples/tutorials/customizing_adanet_with_tfhub.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "fYKSDFyobW3e"
      },
      "source": [
        "From the [customizing AdaNet tutorial](https://github.com/tensorflow/adanet/blob/master/adanet/examples/tutorials/customizing_adanet.ipynb), you know how to define your own neural architecture search space for AdaNet algorithm to explore. One can simplify this process further by using TensorFlow Hub modules as the basic building blocks for AdaNet. These modules have already been pre-trained on large corpuses of data which enables you to leverage the power of transfer learning.\n",
        "\n",
        "In this tutorial, we will create a custom search space for sentiment analysis dataset using TensorFlow Hub text embedding modules.\n",
        "\n"]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "bNRNUJWsbajj"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "# If you are running this in Colab, first install the adanet package:\n",
        "!pip install adanet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "GcA2P5zgbdTT"
      },
      "outputs": [],
      "source": [
        "from __future__ import absolute_import\n",
        "from __future__ import division\n",
        "from __future__ import print_function\n",
        "\n",
        "import functools\n",
        "import os\n",
        "import re\n",
        "import shutil\n",
        "import numpy as np \n",
        "import pandas as pd \n",
        "\n",
        "import tensorflow.compat.v1 as tf\n",
        "import tensorflow_hub as hub\n",
        "\n",
        "import adanet\n",
        "from adanet.examples import simple_dnn\n",
        "\n",
        "# The random seed to use.\n",
        "RANDOM_SEED = 42\n",
        "\n",
        "LOG_DIR = '/tmp/models'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "gc36MZ8Vbrvq"
      },
      "source": [
        "# Getting started\n",
        "\n",
        "## Data\n",
        "We will try to solve the [Large Movie Review Dataset v1.0](http://ai.stanford.edu/~amaas/data/sentiment/) task [(Mass et al., 2011)](http://ai.stanford.edu/~amaas/papers/wvSent_acl2011.pdf). The dataset consists of IMDB movie reviews labeled by positivity from 1 to 10. The task is to label the reviews as **negative** or **positive**."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "4Dx5FxL4bsq5"
      },
      "outputs": [],
      "source": [
        "def load_directory_data(directory):\n",
        "  data = {}\n",
        "  data[\"sentence\"] = []\n",
        "  data[\"sentiment\"] = []\n",
        "  for file_path in os.listdir(directory):\n",
        "    with tf.gfile.GFile(os.path.join(directory, file_path), \"r\") as f:\n",
        "      data[\"sentence\"].append(f.read())\n",
        "      data[\"sentiment\"].append(re.match(\"\\d+_(\\d+)\\.txt\", file_path).group(1))\n",
        "  return pd.DataFrame.from_dict(data)\n",
        "\n",
        "def load_dataset(directory):\n",
        "  pos_df = load_directory_data(os.path.join(directory, \"pos\"))\n",
        "  neg_df = load_directory_data(os.path.join(directory, \"neg\"))\n",
        "  pos_df[\"polarity\"] = 1\n",
        "  neg_df[\"polarity\"] = 0\n",
        "  return pd.concat([pos_df, neg_df]).sample(frac=1).reset_index(drop=True)\n",
        "\n",
        "def download_and_load_datasets(force_download=False):\n",
        "  dataset = tf.keras.utils.get_file(\n",
        "    fname=\"aclImdb.tar.gz\",\n",
        "    origin=\"http://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz\",\n",
        "    extract=True\n",
        "  )\n",
        "  train_df = load_dataset(os.path.join(os.path.dirname(dataset),\n",
        "                                      \"aclImdb\", \"train\"))\n",
        "  test_df = load_dataset(os.path.join(os.path.dirname(dataset),\n",
        "                                      \"aclImdb\", \"test\"))\n",
        "  return train_df, test_df\n",
        "\n",
        "tf.logging.set_verbosity(tf.logging.INFO)\n",
        "\n",
        "train_df, test_df = download_and_load_datasets()\n",
        "train_df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "_N5K2kFsbw4q"
      },
      "source": [
        "## Supply the data in TensorFlow\n",
        "\n",
        "Our first task is to supply the data in TensorFlow. We define three kinds of input_fn that will be used in training later using `pandas_input_fn`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "IbvwjsLTbx4q"
      },
      "outputs": [],
      "source": [
        "FEATURES_KEY = \"sentence\"\n",
        "\n",
        "train_input_fn = tf.estimator.inputs.pandas_input_fn(\n",
        "  train_df, train_df[\"polarity\"], num_epochs=None, shuffle=True)\n",
        "\n",
        "predict_train_input_fn = tf.estimator.inputs.pandas_input_fn(\n",
        "  train_df, train_df[\"polarity\"], shuffle=False)\n",
        "\n",
        "predict_test_input_fn = tf.estimator.inputs.pandas_input_fn(\n",
        "  test_df, test_df[\"polarity\"], shuffle=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "P1YGwRq-b0uo"
      },
      "source": [
        "## Launch TensorBoard\n",
        "\n",
        "Let's run [TensorBoard](https://www.tensorflow.org/guide/summaries_and_tensorboard) to visualize model training over time. We'll use [ngrok](https://ngrok.com/) to tunnel traffic to localhost.\n",
        "\n",
        "*The instructions for setting up Tensorboard were obtained from https://www.dlology.com/blog/quick-guide-to-run-tensorboard-in-google-colab/*\n",
        "\n",
        "Run the next cells and follow the link to see the TensorBoard in a new tab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "K33WZuupb1ja"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "\n",
        "get_ipython().system_raw(\n",
        "    'tensorboard --logdir {} --host 0.0.0.0 --port 6006 \u0026'\n",
        "    .format(LOG_DIR)\n",
        ")\n",
        "\n",
        "# Install ngrok binary.\n",
        "! wget https://bin.equinox.io/c/4VmDzA7iaHb/ngrok-stable-linux-amd64.zip\n",
        "! unzip ngrok-stable-linux-amd64.zip\n",
        "\n",
        "# Delete old logs dir.\n",
        "shutil.rmtree(LOG_DIR, ignore_errors=True)\n",
        "\n",
        "print(\"Follow this link to open TensorBoard in a new tab.\")\n",
        "get_ipython().system_raw('./ngrok http 6006 \u0026')\n",
        "! curl -s http://localhost:4040/api/tunnels | python3 -c \\\n",
        "    \"import sys, json; print(json.load(sys.stdin)['tunnels'][0]['public_url'])\"\n",
        "\n"]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "1BNofTNzb508"
      },
      "source": [
        "## Establish baselines\n",
        "\n",
        "The next task should be to get somes baselines to see how our model performs on\n",
        "this dataset.\n",
        "\n",
        "Let's define some information to share with all our `tf.estimator.Estimators`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "6Y1B8DInb6vZ"
      },
      "outputs": [],
      "source": [
        "NUM_CLASSES = 2\n",
        "\n",
        "loss_reduction = tf.losses.Reduction.SUM_OVER_BATCH_SIZE\n",
        "\n",
        "head = tf.contrib.estimator.binary_classification_head(\n",
        "  loss_reduction=loss_reduction)\n",
        "\n",
        "hub_columns=hub.text_embedding_column(\n",
        "    key=FEATURES_KEY, \n",
        "    module_spec=\"https://tfhub.dev/google/nnlm-en-dim128/1\")\n",
        "\n",
        "def make_config(experiment_name):\n",
        "  # Estimator configuration.\n",
        "  return tf.estimator.RunConfig(\n",
        "    save_checkpoints_steps=1000,\n",
        "    save_summary_steps=1000,\n",
        "    tf_random_seed=RANDOM_SEED,\n",
        "    model_dir=os.path.join(LOG_DIR, experiment_name))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "uX5J5wKPb95j"
      },
      "source": [
        "Let's start simple, and train a linear model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "SRZ31na6b-uK"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "#@title Parameters\n",
        "LEARNING_RATE = 0.001 #@param {type:\"number\"}\n",
        "TRAIN_STEPS = 5000 #@param {type:\"integer\"}\n",
        "\n",
        "estimator = tf.estimator.LinearClassifier(\n",
        "  feature_columns=[hub_columns],\n",
        "  n_classes=NUM_CLASSES,\n",
        "  optimizer=tf.train.RMSPropOptimizer(learning_rate=LEARNING_RATE),\n",
        "  loss_reduction=loss_reduction,\n",
        "  config=make_config(\"linear\"))\n",
        "\n",
        "results, _ = tf.estimator.train_and_evaluate(\n",
        "  estimator,\n",
        "  train_spec=tf.estimator.TrainSpec(\n",
        "    input_fn=train_input_fn,\n",
        "    max_steps=TRAIN_STEPS),\n",
        "  eval_spec=tf.estimator.EvalSpec(\n",
        "    input_fn=predict_test_input_fn,\n",
        "    steps=None))\n",
        "\n",
        "print(\"Accuracy: \", results[\"accuracy\"])\n",
        "print(\"Loss: \", results[\"average_loss\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "wr_oXqQ-cHq_"
      },
      "source": [
        "The linear model with default parameters achieves about **78% accuracy**.\n",
        "\n",
        "Let's see if we can do better with the `simple_dnn` AdaNet:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "3K6hNLc5cIxb"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "#@title Parameters\n",
        "LEARNING_RATE = 0.003  #@param {type:\"number\"}\n",
        "TRAIN_STEPS = 5000  #@param {type:\"integer\"}\n",
        "ADANET_ITERATIONS = 2  #@param {type:\"integer\"}\n",
        "\n",
        "estimator = adanet.Estimator(\n",
        "    head=head,\n",
        "    \n",
        "    # Define the generator, which defines our search space of subnetworks\n",
        "    # to train as candidates to add to the final AdaNet model.\n",
        "    subnetwork_generator=simple_dnn.Generator(\n",
        "        feature_columns=[hub_columns],\n",
        "        optimizer=tf.train.RMSPropOptimizer(learning_rate=LEARNING_RATE),\n",
        "        seed=RANDOM_SEED),\n",
        "    \n",
        "    # The number of train steps per iteration.\n",
        "    max_iteration_steps=TRAIN_STEPS // ADANET_ITERATIONS,\n",
        "    \n",
        "    # The evaluator will evaluate the model on the full training set to\n",
        "    # compute the overall AdaNet loss (train loss + complexity\n",
        "    # regularization) to select the best candidate to include in the\n",
        "    # final AdaNet model.\n",
        "    evaluator=adanet.Evaluator(\n",
        "        input_fn=predict_train_input_fn,\n",
        "        steps=1000),\n",
        "    \n",
        "    # Configuration for Estimators.\n",
        "    config=make_config(\"simple_dnn\"))\n",
        "\n",
        "results, _ = tf.estimator.train_and_evaluate(\n",
        "    estimator,\n",
        "    train_spec=tf.estimator.TrainSpec(\n",
        "        input_fn=train_input_fn,\n",
        "        max_steps=TRAIN_STEPS),\n",
        "    eval_spec=tf.estimator.EvalSpec(\n",
        "        input_fn=predict_test_input_fn,\n",
        "        steps=None))\n",
        "print(\"Accuracy:\", results[\"accuracy\"])\n",
        "print(\"Loss:\", results[\"average_loss\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "-42FKSKBcLsi"
      },
      "source": [
        "The `simple_dnn` AdaNet model with default parameters achieves about **80%\n",
        "accuracy**.\n",
        "\n",
        "This improvement can be attributed to `simple_dnn` searching over\n",
        "fully-connected neural networks which have more expressive power than the linear\n",
        "model due to their non-linear activations.\n",
        "\n",
        "The above `simple_dnn` generator only generates subnetworks that take embedding results from one module. We can add diversity to the search space by building subnetworks that take different embeddings, hence might improve the performance. To do that, we need to define a custom [`adanet.subnetwork.Builder`](https://adanet.readthedocs.io/en/v0.5.0/adanet.subnetwork.html#builder) and [`adanet.subnetwork.Generator`](https://adanet.readthedocs.io/en/v0.5.0/adanet.subnetwork.html#generator)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "arJmpCX7cNvr"
      },
      "source": [
        "## Define a AdaNet model with TensorFlow Hub text embedding modules\n",
        "\n",
        "Creating a new search space for AdaNet to explore is straightforward. There are\n",
        "two abstract classes you need to extend:\n",
        "\n",
        "1.  `adanet.subnetwork.Builder`\n",
        "2.  `adanet.subnetwork.Generator`\n",
        "\n",
        "Similar to the tf.estimator.Estimator `model_fn`, `adanet.subnetwork.Builder`\n",
        "allows you to define your own TensorFlow graph for creating a neural network,\n",
        "and specify the training operations.\n",
        "\n",
        "Below we define one that applies text embedding using TensorFlow Hub text modules first, and then a fully-connected layer to the sentiment polarity."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "R5R5BKjscP2a"
      },
      "outputs": [],
      "source": [
        "class SimpleNetworkBuilder(adanet.subnetwork.Builder):\n",
        "  \"\"\"Builds a simple subnetwork with text embedding module.\"\"\"\n",
        "\n",
        "  def __init__(self, learning_rate, max_iteration_steps, seed,\n",
        "               module_name, module):\n",
        "    \"\"\"Initializes a `SimpleNetworkBuilder`.\n",
        "\n",
        "    Args:\n",
        "      learning_rate: The float learning rate to use.\n",
        "      max_iteration_steps: The number of steps per iteration.\n",
        "      seed: The random seed.\n",
        "\n",
        "    Returns:\n",
        "      An instance of `SimpleNetworkBuilder`.\n",
        "    \"\"\"\n",
        "    self._learning_rate = learning_rate\n",
        "    self._max_iteration_steps = max_iteration_steps\n",
        "    self._seed = seed\n",
        "    self._module_name = module_name\n",
        "    self._module = module\n",
        "\n",
        "  def build_subnetwork(self,\n",
        "                       features,\n",
        "                       logits_dimension,\n",
        "                       training,\n",
        "                       iteration_step,\n",
        "                       summary,\n",
        "                       previous_ensemble=None):\n",
        "    \"\"\"See `adanet.subnetwork.Builder`.\"\"\"\n",
        "    sentence = features[\"sentence\"]\n",
        "    # Load module and apply text embedding, setting trainable=True.\n",
        "    m = hub.Module(self._module, trainable=True)\n",
        "    x = m(sentence)\n",
        "    kernel_initializer = tf.keras.initializers.he_normal(seed=self._seed)\n",
        "\n",
        "    # The `Head` passed to adanet.Estimator will apply the softmax activation.\n",
        "    logits = tf.layers.dense(\n",
        "        x, units=1, activation=None, kernel_initializer=kernel_initializer)\n",
        "\n",
        "    # Use a constant complexity measure, since all subnetworks have the same\n",
        "    # architecture and hyperparameters.\n",
        "    complexity = tf.constant(1)\n",
        "\n",
        "    return adanet.Subnetwork(\n",
        "        last_layer=x,\n",
        "        logits=logits,\n",
        "        complexity=complexity,\n",
        "        persisted_tensors={})\n",
        "\n",
        "  def build_subnetwork_train_op(self, \n",
        "                                subnetwork, \n",
        "                                loss, \n",
        "                                var_list, \n",
        "                                labels, \n",
        "                                iteration_step,\n",
        "                                summary, \n",
        "                                previous_ensemble=None):\n",
        "    \"\"\"See `adanet.subnetwork.Builder`.\"\"\"\n",
        "\n",
        "    learning_rate = tf.train.cosine_decay(\n",
        "        learning_rate=self._learning_rate,\n",
        "        global_step=iteration_step,\n",
        "        decay_steps=self._max_iteration_steps)\n",
        "    optimizer = tf.train.MomentumOptimizer(learning_rate, .9)\n",
        "    # NOTE: The `adanet.Estimator` increments the global step.\n",
        "    return optimizer.minimize(loss=loss, var_list=var_list)\n",
        "\n",
        "  def build_mixture_weights_train_op(self, loss, var_list, logits, labels,\n",
        "                                     iteration_step, summary):\n",
        "    \"\"\"See `adanet.subnetwork.Builder`.\"\"\"\n",
        "    return tf.no_op(\"mixture_weights_train_op\")\n",
        "\n",
        "  @property\n",
        "  def name(self):\n",
        "    \"\"\"See `adanet.subnetwork.Builder`.\"\"\"\n",
        "    return self._module_name"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "oXOcOTlOcSkJ"
      },
      "source": [
        "Next, we extend a `adanet.subnetwork.Generator`, which defines the search\n",
        "space of candidate `SimpleNetworkBuilder` to consider including the final network.\n",
        "It can create one or more at each iteration with different parameters, and the\n",
        "AdaNet algorithm will select the candidate that best improves the overall neural\n",
        "network's `adanet_loss` on the training set.\n",
        "\n",
        "The one below loops through the text embedding modules listed in MODULES and gives it a different random seed at each iteration. These modules are selected from [TensorFlow Hub text modules](https://tfhub.dev/s?module-type=text-embedding):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "TNqiimP-cUcJ"
      },
      "outputs": [],
      "source": [
        "MODULES = [\n",
        "    \"https://tfhub.dev/google/nnlm-en-dim50/1\",\n",
        "    \"https://tfhub.dev/google/nnlm-en-dim128/1\",\n",
        "    \"https://tfhub.dev/google/universal-sentence-encoder/1\"\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "wjqaPzrtcWK3"
      },
      "outputs": [],
      "source": [
        "class SimpleNetworkGenerator(adanet.subnetwork.Generator):\n",
        "  \"\"\"Generates a `SimpleNetwork` at each iteration.\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self, learning_rate, max_iteration_steps, seed=None):\n",
        "    \"\"\"Initializes a `Generator` that builds `SimpleNetwork`.\n",
        "\n",
        "    Args:\n",
        "      learning_rate: The float learning rate to use.\n",
        "      max_iteration_steps: The number of steps per iteration.\n",
        "      seed: The random seed.\n",
        "\n",
        "    Returns:\n",
        "      An instance of `Generator`.\n",
        "    \"\"\"\n",
        "    self._seed = seed\n",
        "    self._dnn_builder_fn = functools.partial(\n",
        "        SimpleNetworkBuilder,\n",
        "        learning_rate=learning_rate,\n",
        "        max_iteration_steps=max_iteration_steps)\n",
        "\n",
        "  def generate_candidates(self, previous_ensemble, iteration_number,\n",
        "                          previous_ensemble_reports, all_reports):\n",
        "    \"\"\"See `adanet.subnetwork.Generator`.\"\"\"\n",
        "    module_index = iteration_number % len(MODULES)\n",
        "    module_name = MODULES[module_index].split(\"/\")[-2]\n",
        "    \n",
        "    print(\"generating candidate: %s\" % module_name)\n",
        "    \n",
        "    seed = self._seed\n",
        "    # Change the seed according to the iteration so that each subnetwork\n",
        "    # learns something different.\n",
        "    if seed is not None:\n",
        "      seed += iteration_number\n",
        "    return [self._dnn_builder_fn(seed=seed, \n",
        "                                 module_name=module_name, \n",
        "                                 module=MODULES[module_index])]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Txpq1eGfcYQb"
      },
      "source": [
        "With these defined, we pass them into a new `adanet.Estimator`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "vWI3laQscaMn"
      },
      "outputs": [],
      "source": [
        "#@title Parameters\n",
        "LEARNING_RATE = 0.05  #@param {type:\"number\"}\n",
        "TRAIN_STEPS = 7500  #@param {type:\"integer\"}\n",
        "ADANET_ITERATIONS = 3  #@param {type:\"integer\"}\n",
        "\n",
        "max_iteration_steps = TRAIN_STEPS // ADANET_ITERATIONS\n",
        "estimator = adanet.Estimator(\n",
        "    head=head,\n",
        "    subnetwork_generator=SimpleNetworkGenerator(\n",
        "        learning_rate=LEARNING_RATE,\n",
        "        max_iteration_steps=max_iteration_steps,\n",
        "        seed=RANDOM_SEED),\n",
        "    max_iteration_steps=max_iteration_steps,\n",
        "    evaluator=adanet.Evaluator(input_fn=train_input_fn, \n",
        "                               steps=10),\n",
        "    report_materializer=None,\n",
        "    adanet_loss_decay=.99,\n",
        "    config=make_config(\"tfhub\"))\n",
        "\n",
        "results, _ = tf.estimator.train_and_evaluate(\n",
        "    estimator,\n",
        "    train_spec=tf.estimator.TrainSpec(input_fn=train_input_fn,\n",
        "                                      max_steps=TRAIN_STEPS),\n",
        "    eval_spec=tf.estimator.EvalSpec(input_fn=predict_test_input_fn, \n",
        "                                    steps=None))\n",
        "print(\"Accuracy:\", results[\"accuracy\"])\n",
        "print(\"Loss:\", results[\"average_loss\"])\n",
        "\n",
        "\n",
        "\n",
        "def ensemble_architecture(result):\n",
        "  \"\"\"Extracts the ensemble architecture from evaluation results.\"\"\"\n",
        "\n",
        "  architecture = result[\"architecture/adanet/ensembles\"]\n",
        "  # The architecture is a serialized Summary proto for TensorBoard.\n",
        "  summary_proto = tf.summary.Summary.FromString(architecture)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "eZf_k22_cdki"
      },
      "source": [
        "Our SimpleNetworkGenerator code achieves about \u003cb\u003e87% accuracy \u003c/b\u003e, which is almost \u003cb\u003e7%\u003c/b\u003e higher than with using just one network directly.\n",
        "\n",
        "You can see how the performance improves step by step:\n",
        "\n",
        "| Linear Baseline | Adanet + simple_dnn | Adanet + TensorFlow Hub  |\n",
        "| --- |:---:| ---:|\n",
        "| 78% | 80%| 87% |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "O67LjSUTceRN"
      },
      "source": [
        "## Generating predictions on our trained model\n",
        "\n",
        "Now that we've got a trained model, we can use it to generate predictions on new input. To keep things simple, here we'll generate predictions on our `estimator` using the first 10 examples from the test set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "LMXCcabVcgZu"
      },
      "outputs": [],
      "source": [
        "predict_input_fn = tf.estimator.inputs.pandas_input_fn(\n",
        "  test_df.iloc[:10], test_df[\"polarity\"].iloc[:10], shuffle=False)\n",
        "\n",
        "predictions = estimator.predict(input_fn=predict_input_fn)\n",
        "\n",
        "for i, val in enumerate(predictions):\n",
        "    predicted_class = val['class_ids'][0]\n",
        "    prediction_confidence = val['probabilities'][predicted_class] * 100\n",
        "    \n",
        "    print('Actual text: ' + test_df[\"sentence\"][i])\n",
        "    print('Predicted class: %s, confidence: %s%%' \n",
        "          % (predicted_class, round(prediction_confidence, 3)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "RnMGQnGkcnVY"
      },
      "source": [
        "## Conclusion and next steps\n",
        "\n",
        "In this tutorial, you learned how to customize `adanet` to encode your\n",
        "understanding of a particular dataset, and explore novel search spaces with\n",
        "AdaNet with TensorFlow Hub modules.\n",
        "\n",
        "As an exercise, you can swap out the ACL IMDB dataset with other text dataset in this notebook and see how `SimpleNetworkGenerator` performs."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
            "name": "customizing_adanet_with_tfhub.ipynb",
      "provenance": [],
      "version": "0.3.2"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
